"""
This script is used to connect to the MCP server and pipe the input and output to multiple websocket endpoints.
Version: 0.2.0

Usage:

export MCP_ENDPOINT=<mcp_endpoint>
python mcp_pipe_multi.py <mcp_script1> <mcp_script2> ...
"""

import asyncio
import websockets
import subprocess
import logging
import os
import signal
import sys
import random
from dotenv import load_dotenv
from typing import Dict, List, Optional

# Load environment variables from .env file
load_dotenv()

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('MCP_PIPE_MULTI')

# Reconnection settings
INITIAL_BACKOFF = 1  # Initial wait time in seconds
MAX_BACKOFF = 600  # Maximum wait time in seconds

class ToolProcess:
    """Class to manage a single tool process and its WebSocket connection"""
    def __init__(self, script_path: str, endpoint_url: str):
        self.script_path = script_path
        self.endpoint_url = endpoint_url
        self.process: Optional[subprocess.Popen] = None
        self.websocket: Optional[websockets.WebSocketClientProtocol] = None
        self.reconnect_attempt = 0
        self.backoff = INITIAL_BACKOFF
        self.running = True
    
    async def connect_with_retry(self):
        """Connect to WebSocket server with retry mechanism"""
        while self.running:
            try:
                if self.reconnect_attempt > 0:
                    wait_time = self.backoff * (1 + random.random() * 0.1)  # Add some random jitter
                    logger.info(f"[{self.script_path}] Waiting {wait_time:.2f} seconds before reconnection attempt {self.reconnect_attempt}...")
                    await asyncio.sleep(wait_time)
                    
                # Attempt to connect
                await self.connect_to_server()
            except Exception as e:
                self.reconnect_attempt += 1
                logger.warning(f"[{self.script_path}] Connection closed (attempt: {self.reconnect_attempt}): {e}")
                # Calculate wait time for next reconnection (exponential backoff)
                self.backoff = min(self.backoff * 2, MAX_BACKOFF)
            
            # Ensure the process is terminated if connection failed
            if self.process:
                self.terminate_process()
    
    async def connect_to_server(self):
        """Connect to WebSocket server and establish bidirectional communication"""
        try:
            logger.info(f"[{self.script_path}] Connecting to WebSocket server...")
            async with websockets.connect(self.endpoint_url) as websocket:
                self.websocket = websocket
                logger.info(f"[{self.script_path}] Successfully connected to WebSocket server")
                
                # Reset reconnection counter
                self.reconnect_attempt = 0
                self.backoff = INITIAL_BACKOFF
                
                # Start tool process
                self.process = subprocess.Popen(
                    ['python', self.script_path],
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    encoding='utf-8',
                    text=True
                )
                logger.info(f"[{self.script_path}] Started process")
                
                # Create tasks for bidirectional communication
                await asyncio.gather(
                    self.pipe_websocket_to_process(),
                    self.pipe_process_to_websocket(),
                    self.pipe_process_stderr_to_terminal()
                )
        except websockets.exceptions.ConnectionClosed as e:
            logger.error(f"[{self.script_path}] WebSocket connection closed: {e}")
            raise
        except Exception as e:
            logger.error(f"[{self.script_path}] Connection error: {e}")
            raise
    
    async def pipe_websocket_to_process(self):
        """Read data from WebSocket and write to process stdin"""
        if not self.process or not self.websocket:
            return
            
        try:
            while True:
                # Read message from WebSocket
                message = await self.websocket.recv()
                logger.debug(f"[{self.script_path}] << {message[:120]}...")
                
                # Write to process stdin
                if isinstance(message, bytes):
                    message = message.decode('utf-8')
                self.process.stdin.write(message + '\n')
                self.process.stdin.flush()
        except Exception as e:
            logger.error(f"[{self.script_path}] Error in WebSocket to process pipe: {e}")
            raise
        finally:
            # Close process stdin
            if self.process and not self.process.stdin.closed:
                self.process.stdin.close()
    
    async def pipe_process_to_websocket(self):
        """Read data from process stdout and send to WebSocket"""
        if not self.process or not self.websocket:
            return
            
        try:
            while True:
                # Read data from process stdout
                data = await asyncio.get_event_loop().run_in_executor(
                    None, self.process.stdout.readline
                )
                
                if not data:
                    logger.info(f"[{self.script_path}] Process has ended output")
                    break
                    
                # Send data to WebSocket
                logger.debug(f"[{self.script_path}] >> {data[:120]}...")
                await self.websocket.send(data)
        except Exception as e:
            logger.error(f"[{self.script_path}] Error in process to WebSocket pipe: {e}")
            raise
    
    async def pipe_process_stderr_to_terminal(self):
        """Read data from process stderr and print to terminal"""
        if not self.process:
            return
            
        try:
            while True:
                # Read data from process stderr
                data = await asyncio.get_event_loop().run_in_executor(
                    None, self.process.stderr.readline
                )
                
                if not data:
                    logger.info(f"[{self.script_path}] Process has ended stderr output")
                    break
                    
                # Print stderr data to terminal
                sys.stderr.write(f"[{self.script_path}] {data}")
                sys.stderr.flush()
        except Exception as e:
            logger.error(f"[{self.script_path}] Error in process stderr pipe: {e}")
            raise
    
    def terminate_process(self):
        """Terminate the tool process"""
        if self.process:
            logger.info(f"[{self.script_path}] Terminating process")
            try:
                self.process.terminate()
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.process.kill()
            logger.info(f"[{self.script_path}] Process terminated")
            self.process = None
    
    def stop(self):
        """Stop the tool process and connection"""
        self.running = False
        self.terminate_process()
        if self.websocket:
            asyncio.create_task(self.websocket.close())

async def main(scripts: List[str], endpoint_url: str):
    """Main function to start multiple tool processes"""
    tool_processes = [ToolProcess(script, endpoint_url) for script in scripts]
    
    try:
        # Start all tool processes
        tasks = [process.connect_with_retry() for process in tool_processes]
        await asyncio.gather(*tasks)
    except asyncio.CancelledError:
        logger.info("Main task cancelled")
    finally:
        # Clean up all processes
        for process in tool_processes:
            process.stop()

if __name__ == "__main__":
    # Register signal handler
    signal.signal(signal.SIGINT, lambda sig, frame: sys.exit(0))
    
    # Check for script arguments
    if len(sys.argv) < 2:
        logger.error("Usage: mcp_pipe_multi.py <mcp_script1> <mcp_script2> ...")
        sys.exit(1)
    
    script_paths = sys.argv[1:]
    
    # Get endpoint URL from environment variable
    endpoint_url = os.environ.get('MCP_ENDPOINT')
    if not endpoint_url:
        logger.error("Please set the `MCP_ENDPOINT` environment variable")
        sys.exit(1)
    
    # Start main loop
    try:
        asyncio.run(main(script_paths, endpoint_url))
    except KeyboardInterrupt:
        logger.info("Program interrupted by user")
    except Exception as e:
        logger.error(f"Program execution error: {e}")